"""
车牌检测模块
使用YOLO模型检测图像中的车牌位置
"""

import os
import cv2
import numpy as np
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from ultralytics import YOLO
from src.config import config
from src.utils.logger import app_logger
from src.utils.image_utils import image_processor


class PlateDetector:
    """
    车牌检测器类
    
    用途：
    - 使用YOLO模型检测图像中的车牌
    - 提供车牌位置的边界框坐标
    - 支持批量检测和单张图像检测
    
    使用场景：
    - 从监控图像中检测车牌
    - 为车牌识别提供准确的车牌区域
    - 车辆管理系统中的车牌定位
    """
    
    def __init__(self, model_path: str = None):
        """
        初始化车牌检测器
        
        参数：
            model_path (str): YOLO模型文件路径，默认使用配置中的模型
        """
        self.model_path = model_path or config.get('model.yolo_model')
        self.confidence_threshold = config.get('model.confidence_threshold')
        self.iou_threshold = config.get('model.iou_threshold')
        self.model = None
        
        self._initialize_model()
    
    def _initialize_model(self):
        """
        初始化YOLO模型
        """
        try:
            # 确保模型目录存在
            models_dir = Path(config.get('paths.models_dir'))
            models_dir.mkdir(parents=True, exist_ok=True)
            
            # 构建完整的模型路径
            full_model_path = models_dir / self.model_path
            
            if full_model_path.exists():
                # 模型文件存在，直接加载
                app_logger.info(f"加载本地模型: {full_model_path}")
                self.model = YOLO(str(full_model_path))
            else:
                # 模型文件不存在，需要下载
                app_logger.info(f"模型文件不存在，开始下载: {self.model_path}")
                
                # 下载模型到临时位置
                temp_model = YOLO(self.model_path)
                
                # 尝试从多个可能的缓存位置找到下载的模型文件
                import shutil
                import os
                
                possible_cache_paths = [
                    Path.home() / '.cache' / 'ultralytics' / 'yolov8n.pt',
                    Path.home() / '.ultralytics' / 'yolov8n.pt',
                    Path.home() / '.cache' / 'torch' / 'hub' / 'ultralytics_yolov5_master' / 'yolov8n.pt',
                ]
                
                # 添加当前工作目录（有时模型会下载到当前目录）
                if (Path.cwd() / self.model_path).exists():
                    possible_cache_paths.insert(0, Path.cwd() / self.model_path)
                
                model_copied = False
                for cache_path in possible_cache_paths:
                    if cache_path.exists():
                        app_logger.info(f"找到缓存模型: {cache_path}")
                        shutil.copy2(cache_path, full_model_path)
                        app_logger.info(f"模型已复制到: {full_model_path}")
                        model_copied = True
                        
                        # 删除原始缓存文件以节省空间（可选）
                        try:
                            if cache_path != full_model_path:  # 确保不删除我们刚复制的文件
                                cache_path.unlink()
                                app_logger.info(f"已删除缓存文件: {cache_path}")
                        except Exception as e:
                            app_logger.warning(f"删除缓存文件失败: {e}")
                        
                        break
                
                if model_copied:
                    # 加载复制后的模型
                    self.model = YOLO(str(full_model_path))
                    app_logger.info(f"成功加载模型: {full_model_path}")
                else:
                    # 如果无法找到缓存文件，使用临时模型
                    app_logger.warning("无法找到缓存模型文件，使用临时模型")
                    self.model = temp_model
            
            # 验证模型是否加载成功
            if self.model is None:
                raise RuntimeError("模型加载失败")
            
            # 记录模型信息
            app_logger.info("车牌检测模型初始化成功")
            if hasattr(self.model, 'model'):
                app_logger.info(f"模型类型: {type(self.model.model).__name__}")
            
        except Exception as e:
            app_logger.error(f"初始化YOLO模型失败: {e}")
            raise
    
    def detect_plates(self, image: np.ndarray) -> List[Dict]:
        """
        检测图像中的车牌
        
        参数：
            image (np.ndarray): 输入图像，BGR格式
            
        返回值：
            List[Dict]: 检测结果列表，每个元素包含：
                - bbox: 边界框坐标 (x1, y1, x2, y2)
                - confidence: 置信度
                - class_id: 类别ID
                - class_name: 类别名称
        """
        if self.model is None:
            app_logger.error("模型未初始化")
            return []
        
        try:
            # 预处理图像
            processed_image = image_processor.preprocess_for_detection(image)
            
            # 使用YOLO进行检测，添加多尺度检测
            results = self.model(
                processed_image,
                conf=self.confidence_threshold,
                iou=self.iou_threshold,
                verbose=False,
                imgsz=640,  # 标准尺寸
                augment=True  # 启用测试时增强
            )
            
            detections = []
            
            # 处理检测结果
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for box in boxes:
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        
                        # 获取置信度
                        confidence = box.conf[0].cpu().numpy()
                        
                        # 获取类别信息
                        class_id = int(box.cls[0].cpu().numpy())
                        class_name = self.model.names[class_id] if hasattr(self.model, 'names') else str(class_id)
                        
                        # 转换坐标为整数
                        bbox = (int(x1), int(y1), int(x2), int(y2))
                        
                        detection = {
                            'bbox': bbox,
                            'confidence': float(confidence),
                            'class_id': class_id,
                            'class_name': class_name
                        }
                        
                        detections.append(detection)
            
            app_logger.info(f"检测到 {len(detections)} 个车牌")
            return detections
            
        except Exception as e:
            app_logger.error(f"车牌检测失败: {e}")
            return []
    
    def detect_plates_from_file(self, image_path: str) -> List[Dict]:
        """
        从文件检测车牌
        
        参数：
            image_path (str): 图像文件路径
            
        返回值：
            List[Dict]: 检测结果列表
        """
        # 加载图像
        image = image_processor.load_image(image_path)
        if image is None:
            return []
        
        # 执行检测
        return self.detect_plates(image)
    
    def extract_plate_images(self, image: np.ndarray, detections: List[Dict]) -> List[Dict]:
        """
        从检测结果中提取车牌图像
        
        参数：
            image (np.ndarray): 原始图像
            detections (List[Dict]): 检测结果列表
            
        返回值：
            List[Dict]: 包含车牌图像的检测结果列表
        """
        results = []
        
        for i, detection in enumerate(detections):
            bbox = detection['bbox']
            
            # 裁剪车牌图像
            plate_image = image_processor.crop_plate(image, bbox)
            
            if plate_image is not None:
                # 复制检测结果并添加车牌图像
                result = detection.copy()
                result['plate_image'] = plate_image
                result['plate_id'] = i
                results.append(result)
            else:
                app_logger.warning(f"无法提取车牌图像，索引: {i}")
        
        return results
    
    def filter_detections(self, detections: List[Dict], min_confidence: float = None) -> List[Dict]:
        """
        过滤检测结果
        
        参数：
            detections (List[Dict]): 检测结果列表
            min_confidence (float): 最小置信度阈值
            
        返回值：
            List[Dict]: 过滤后的检测结果列表
        """
        if min_confidence is None:
            min_confidence = self.confidence_threshold
        
        filtered = [det for det in detections if det['confidence'] >= min_confidence]
        
        app_logger.info(f"过滤后保留 {len(filtered)} 个检测结果")
        return filtered
    
    def get_best_detection(self, detections: List[Dict]) -> Optional[Dict]:
        """
        获取置信度最高的检测结果
        
        参数：
            detections (List[Dict]): 检测结果列表
            
        返回值：
            Dict: 最佳检测结果，如果没有检测结果则返回None
        """
        if not detections:
            return None
        
        best_detection = max(detections, key=lambda x: x['confidence'])
        return best_detection
    
    def update_model(self, new_model_path: str):
        """
        更新检测模型
        
        参数：
            new_model_path (str): 新模型文件路径
        """
        try:
            self.model_path = new_model_path
            self._initialize_model()
            app_logger.info(f"模型更新成功: {new_model_path}")
        except Exception as e:
            app_logger.error(f"模型更新失败: {e}")
            raise
    
    def set_confidence_threshold(self, threshold: float):
        """
        设置置信度阈值
        
        参数：
            threshold (float): 置信度阈值 (0.0-1.0)
        """
        if 0.0 <= threshold <= 1.0:
            self.confidence_threshold = threshold
            app_logger.info(f"置信度阈值设置为: {threshold}")
        else:
            app_logger.error(f"无效的置信度阈值: {threshold}")
    
    def set_iou_threshold(self, threshold: float):
        """
        设置IoU阈值
        
        参数：
            threshold (float): IoU阈值 (0.0-1.0)
        """
        if 0.0 <= threshold <= 1.0:
            self.iou_threshold = threshold
            app_logger.info(f"IoU阈值设置为: {threshold}")
        else:
            app_logger.error(f"无效的IoU阈值: {threshold}")
    
    def get_model_info(self) -> Dict:
        """
        获取模型信息
        
        返回值：
            Dict: 模型信息字典
        """
        models_dir = Path(config.get('paths.models_dir'))
        full_model_path = models_dir / self.model_path
        
        return {
            'model_name': self.model_path,
            'model_path': str(full_model_path),
            'model_exists': full_model_path.exists(),
            'model_size': full_model_path.stat().st_size if full_model_path.exists() else 0,
            'confidence_threshold': self.confidence_threshold,
            'iou_threshold': self.iou_threshold,
            'model_loaded': self.model is not None
        }


# 全局车牌检测器实例
plate_detector = PlateDetector() 